import torch
from torch import nn


class ResidualBlock(nn.Module):
    def __init__(self, in_channel, out_channel):
        super().__init__()
        self.bn1 = nn.Sequential(
            nn.Conv2d(in_channel, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )

        self.bn2 = nn.Sequential(
            nn.Conv2d(64, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, out_channel, 3, 1, 1),
            nn.BatchNorm2d(out_channel)
        )

        self.align = nn.Conv2d(in_channel, out_channel, 3, 1, 1)

        self.fn = nn.ReLU()


    def forward(self, x):

        align = self.align(x)

        x = self.bn1(x)
        x = self.bn2(x)

        return self.fn(align + x)
    

class ResidualBottleBlock(nn.Module):
    def __init__(self, in_channel, out_channel):
        super().__init__()

        self.bn1 = nn.Sequential(
            nn.Conv2d(in_channel, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )

        self.bn2 = nn.Sequential(
            nn.Conv2d(64, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, out_channel, 3, 1, 1),
            nn.BatchNorm2d(out_channel)
        )

        self.align = nn.Conv2d(in_channel, out_channel, 3, 1, 1)

        self.fn = nn.ReLU()

    def forward(self, x):

        align = self.align(x)

        x = self.bn1(x)
        x = self.bn2(x)

        return self.fn(align + x)


class VAEEncoder(nn.Module):
    def __init__(self, channel, z_dim):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channel, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.AvgPool2d(2, 2)
        )

        self.b1 = nn.Sequential(
            nn.Conv2d(64, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AvgPool2d(2, 2)
        )

        self.b2 = nn.Sequential(
            nn.Conv2d(256, 256, 3, 1, 1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.AvgPool2d(2, 2)
        )

        self.flatten = nn.Flatten()

        self.align = nn.Linear(64*3*3, z_dim)


    def forward(self, x):
        x = self.conv(x)
        x = self.b1(x)
        x = self.b2(x)
        x = self.flatten(x)

        return self.align(x)
    


class Block(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()

        self.space = nn.Sequential(
            nn.Linear(in_dim, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Linear(64, out_dim),
            nn.BatchNorm1d(out_dim)
        )

        self.align = nn.Linear(in_dim, out_dim)

        self.fn = nn.ReLU()

    def forward(self, x):
        return self.fn(self.space(x) + self.align(x))

class VAEDecoder(nn.Module):
    def __init__(self, z_dim, out_dim):
        super().__init__()

        self.block1 = nn.Sequential(
            nn.Linear(z_dim, 1024),
            nn.BatchNorm1d(1024),
            nn.ReLU()
        )

        self.block2 = nn.Linear(1024, out_dim)

        self.std = nn.Sigmoid()

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        return self.std(x)

class VAE(nn.Module):
    def __init__(self, channel, latent_dim, z_dim, out_dim):
        super().__init__()

        self.z_dim = z_dim
        self.encoder = VAEEncoder(channel, latent_dim)

        self.var = nn.Linear(latent_dim, z_dim)
        self.mean = nn.Linear(latent_dim,z_dim)

        self.decoder = VAEDecoder(z_dim, out_dim)


    def forward(self, x):
        latent = self.encoder(x)
        var = self.var(latent)
        mean = self.mean(latent)

        eps = torch.normal(0, 1, size=var.size()).to(x.device)

        sample = mean + eps*var

        return mean, var, self.decoder(sample)
    

class VAELoss(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.mse = nn.MSELoss()

    def forward(self, mean, var, img, target):
        l1 = self.mse(img, target)
        l2 = torch.sum(mean**2 + torch.exp(var) - var - 1) / mean.size(0)

        return l1 + l2